package org.summerframework.component.security.session.websocket.interceptor;

import org.summerframework.component.security.LoginData;
import org.springframework.messaging.Message;
import org.springframework.messaging.simp.SimpMessageHeaderAccessor;
import org.springframework.messaging.simp.SimpMessageType;
import org.springframework.messaging.simp.stomp.StompHeaderAccessor;
import org.springframework.messaging.support.MessageHeaderAccessor;
import org.springframework.messaging.support.NativeMessageHeaderAccessor;
import org.springframework.session.Session;
import org.springframework.session.SessionRepository;

import java.security.Principal;
import java.time.Instant;
import java.util.HashMap;
import java.util.Map;

/**
 * 给webSocket 添加用户
 * {@link LoginData} 继承于 {@link Principal}
 *
 * @param <S> httpSession
 */
public abstract class AbstractHandshakeInterceptor<S extends Session> {

    final SessionRepository<S> sessionRepository;

    static final String SPRING_SESSION_ID_ATTR_NAME = "SPRING.SESSION.ID";
    static final String SPRING_SESSION = "SPRING.SESSION";

    public AbstractHandshakeInterceptor(SessionRepository<S> sessionRepository) {
        this.sessionRepository = sessionRepository;
    }

    public S getSession(Message<?> message) {
        SimpMessageHeaderAccessor simpMessageHeaderAccessor = SimpMessageHeaderAccessor.wrap(message);

        Object header = simpMessageHeaderAccessor.getHeader(SPRING_SESSION);
        if (header != null) {
            return (S) header;
        }

        Map<String, Object> sessionHeaders = simpMessageHeaderAccessor.getSessionAttributes();

        if (sessionHeaders == null) {
            sessionHeaders = new HashMap<>(16);
            simpMessageHeaderAccessor.setSessionAttributes(sessionHeaders);
        }

        if (SimpMessageHeaderAccessor.getMessageType(message.getHeaders()) == SimpMessageType.CONNECT) {
            NativeMessageHeaderAccessor nativeMessageHeaderAccessor = MessageHeaderAccessor.getAccessor(message, NativeMessageHeaderAccessor.class);

            if (nativeMessageHeaderAccessor != null) {
                String sessionId = nativeMessageHeaderAccessor.getFirstNativeHeader("sessionId");
                if (sessionId != null) {
                    setSessionId(sessionHeaders, sessionId);
                }
            }
        }

        String sessionId = getSessionId(sessionHeaders);
        String receipt = simpMessageHeaderAccessor.getFirstNativeHeader(StompHeaderAccessor.STOMP_RECEIPT_HEADER);

        if (receipt != null) {
            sessionHeaders.put(StompHeaderAccessor.STOMP_RECEIPT_HEADER, receipt);
        }

        if (sessionId != null) {
            S session = this.sessionRepository.findById(sessionId);

            if (session != null) {
                session.setLastAccessedTime(Instant.now());
                simpMessageHeaderAccessor.setHeader(SPRING_SESSION, session);
                sessionRepository.save(session);
            }
            return session;
        } else {
            return null;
        }
    }

    public static String getSessionId(Map<String, Object> attributes) {
        return (String) attributes.get(SPRING_SESSION_ID_ATTR_NAME);
    }

    public static void setSessionId(Map<String, Object> attributes, String sessionId) {
        attributes.put(SPRING_SESSION_ID_ATTR_NAME, sessionId);
    }
}
